In this tutorial we solve the optimal control problem
$$\min J(y, u) = \frac{1}{2} \int_{\Omega} |v - v_d|^2 dx + \frac{\alpha}{2} \int_{\Omega} |u|^2 dx$$s.t. $$\begin{cases}
\end{cases}$$
where $$\begin{align*} & \Omega & \text{unit square}\\ & u \in [L^2(\Omega)]^2 & \text{control variable}\\ & v \in [H^1_0(\Omega)]^2 & \text{state velocity variable}\\ & p \in L^2(\Omega) & \text{state pressure variable}\\ & \alpha > 0 & \text{penalization parameter}\\ & v_d & \text{desired state}\\ & f & \text{forcing term} \end{align*}$$ using an adjoint formulation solved by a one shot approach
import dolfinx.fem
import dolfinx.io
import dolfinx.mesh
import mpi4py.MPI
import numpy as np
import numpy.typing
import petsc4py.PETSc
import sympy
import ufl
import viskex
import multiphenicsx.fem
mesh = dolfinx.mesh.create_unit_square(mpi4py.MPI.COMM_WORLD, 32, 32)
def bottom(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
"""Condition that defines the bottom boundary."""
return abs(x[1] - 0.) < np.finfo(float).eps # type: ignore[no-any-return]
def left(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
"""Condition that defines the left boundary."""
return abs(x[0] - 0.) < np.finfo(float).eps # type: ignore[no-any-return]
def top(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
"""Condition that defines the top boundary."""
return abs(x[1] - 1.) < np.finfo(float).eps # type: ignore[no-any-return]
def right(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
"""Condition that defines the right boundary."""
return abs(x[0] - 1.) < np.finfo(float).eps # type: ignore[no-any-return]
boundaries_entities = dict()
boundaries_values = dict()
for (boundary, boundary_id) in zip((bottom, left, top, right), (1, 2, 3, 4)):
boundaries_entities[boundary_id] = dolfinx.mesh.locate_entities_boundary(
mesh, mesh.topology.dim - 1, boundary)
boundaries_values[boundary_id] = np.full(
boundaries_entities[boundary_id].shape, boundary_id, dtype=np.int32)
boundaries_entities_unsorted = np.hstack(list(boundaries_entities.values()))
boundaries_values_unsorted = np.hstack(list(boundaries_values.values()))
boundaries_entities_argsort = np.argsort(boundaries_entities_unsorted)
boundaries_entities_sorted = boundaries_entities_unsorted[boundaries_entities_argsort]
boundaries_values_sorted = boundaries_values_unsorted[boundaries_entities_argsort]
boundaries = dolfinx.mesh.meshtags(
mesh, mesh.topology.dim - 1,
boundaries_entities_sorted, boundaries_values_sorted)
boundaries.name = "boundaries"
boundaries_1234 = boundaries.indices[np.isin(boundaries.values, (1, 2, 3, 4))]
viskex.dolfinx.plot_mesh(mesh)
viskex.dolfinx.plot_mesh_entities(mesh, mesh.topology.dim - 1, "boundaries", boundaries_1234)
Y_velocity = dolfinx.fem.VectorFunctionSpace(mesh, ("Lagrange", 2))
Y_pressure = dolfinx.fem.FunctionSpace(mesh, ("Lagrange", 1))
U = dolfinx.fem.VectorFunctionSpace(mesh, ("Lagrange", 2))
Q_velocity = Y_velocity.clone()
Q_pressure = Y_pressure.clone()
(v, p) = (ufl.TrialFunction(Y_velocity), ufl.TrialFunction(Y_pressure))
(w, q) = (ufl.TestFunction(Y_velocity), ufl.TestFunction(Y_pressure))
u = ufl.TrialFunction(U)
r = ufl.TestFunction(U)
(z, b) = (ufl.TrialFunction(Q_velocity), ufl.TrialFunction(Q_pressure))
(s, d) = (ufl.TestFunction(Q_velocity), ufl.TestFunction(Q_pressure))
alpha = 1.e-5
epsilon = 1.e-5
x, y = sympy.symbols("x[0], x[1]")
psi_d = 10 * (1 - sympy.cos(0.8 * np.pi * x)) * (1 - sympy.cos(0.8 * np.pi * y)) * (1 - x)**2 * (1 - y)**2
v_d_x = sympy.lambdify([x, y], psi_d.diff(y, 1))
v_d_y = sympy.lambdify([x, y], - psi_d.diff(x, 1))
v_d = dolfinx.fem.Function(Y_velocity)
v_d.interpolate(lambda x: np.stack((v_d_x(x[0], x[1]), v_d_y(x[0], x[1])), axis=0))
ff = dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2)))
bc0 = np.zeros((2, ), dtype=petsc4py.PETSc.ScalarType)
a = [[ufl.inner(v, w) * ufl.dx, None, None, ufl.inner(ufl.grad(z), ufl.grad(w)) * ufl.dx,
- ufl.inner(b, ufl.div(w)) * ufl.dx],
[None, None, None, - ufl.inner(ufl.div(z), q) * ufl.dx, epsilon * ufl.inner(b, q) * ufl.dx],
[None, None, alpha * ufl.inner(u, r) * ufl.dx, - ufl.inner(z, r) * ufl.dx, None],
[ufl.inner(ufl.grad(v), ufl.grad(s)) * ufl.dx, - ufl.inner(p, ufl.div(s)) * ufl.dx,
- ufl.inner(u, s) * ufl.dx, None, None],
[- ufl.inner(ufl.div(v), d) * ufl.dx, epsilon * ufl.inner(p, d) * ufl.dx, None, None, None]]
f = [ufl.inner(v_d, w) * ufl.dx,
None,
None,
ufl.inner(ff, s) * ufl.dx,
None]
a[3][3] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(z, s) * ufl.dx
f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), q) * ufl.dx
f[2] = ufl.inner(dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2))), r) * ufl.dx
f[4] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), d) * ufl.dx
a_cpp = dolfinx.fem.form(a)
f_cpp = dolfinx.fem.form(f)
bdofs_Y_velocity_1234 = dolfinx.fem.locate_dofs_topological(
Y_velocity, mesh.topology.dim - 1, boundaries_1234)
bdofs_Q_velocity_1234 = dolfinx.fem.locate_dofs_topological(
Q_velocity, mesh.topology.dim - 1, boundaries_1234)
bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_velocity_1234, Y_velocity),
dolfinx.fem.dirichletbc(bc0, bdofs_Q_velocity_1234, Q_velocity)]
(v, p) = (dolfinx.fem.Function(Y_velocity), dolfinx.fem.Function(Y_pressure))
u = dolfinx.fem.Function(U)
(z, b) = (dolfinx.fem.Function(Q_velocity), dolfinx.fem.Function(Q_pressure))
J = 0.5 * ufl.inner(v - v_d, v - v_d) * ufl.dx + 0.5 * alpha * ufl.inner(u, u) * ufl.dx
J_cpp = dolfinx.fem.form(J)
# Extract state forms from the optimality conditions
a_state = [[ufl.replace(a[i][j], {s: w, d: q}) if a[i][j] is not None else None
for j in (0, 1)] for i in (3, 4)]
f_state = [ufl.replace(f[i], {s: w, d: q}) for i in (3, 4)]
a_state_cpp = dolfinx.fem.form(a_state)
f_state_cpp = dolfinx.fem.form(f_state)
bc_state = [bc[0]]
# Assemble the block linear system for the state
A_state = dolfinx.fem.petsc.assemble_matrix_block(a_state_cpp, bcs=bc_state)
A_state.assemble()
F_state = dolfinx.fem.petsc.assemble_vector_block(f_state_cpp, a_state_cpp, bcs=bc_state)
# Solve
vp = dolfinx.fem.petsc.create_vector_block([f_cpp[j] for j in (0, 1)])
ksp = petsc4py.PETSc.KSP()
ksp.create(mesh.comm)
ksp.setOperators(A_state)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(F_state, vp)
vp.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()
<petsc4py.PETSc.KSP at 0x7fca8d755d00>
# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(vp, [c.function_space.dofmap for c in (v, p)]) as vp_wrapper:
for vp_wrapper_local, component in zip(vp_wrapper, (v, p)):
with component.vector.localForm() as component_local:
component_local[:] = vp_wrapper_local
J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)
print("Uncontrolled J =", J_uncontrolled)
assert np.isclose(J_uncontrolled, 0.1784536)
Uncontrolled J = 0.17845361493420536
viskex.dolfinx.plot_vector_field(v, "uncontrolled state velocity", glyph_factor=1e-1)
viskex.dolfinx.plot_scalar_field(p, "uncontrolled state pressure")
# Assemble the block linear system for the optimality conditions
A = dolfinx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc)
A.assemble()
F = dolfinx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc)
# Solve
vpuzb = dolfinx.fem.petsc.create_vector_block(f_cpp)
ksp = petsc4py.PETSc.KSP()
ksp.create(mesh.comm)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(F, vpuzb)
vpuzb.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()
<petsc4py.PETSc.KSP at 0x7fca8d78c770>
# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(
vpuzb, [c.function_space.dofmap for c in (v, p, u, z, b)]) as vpuzb_wrapper:
for vpuzb_wrapper_local, component in zip(vpuzb_wrapper, (v, p, u, z, b)):
with component.vector.localForm() as component_local:
component_local[:] = vpuzb_wrapper_local
J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)
print("Optimal J =", J_controlled)
assert np.isclose(J_controlled, 0.0052941)
Optimal J = 0.005294106310254608
viskex.dolfinx.plot_vector_field(v, "state velocity", glyph_factor=1e-1)
viskex.dolfinx.plot_scalar_field(p, "state pressure")
viskex.dolfinx.plot_vector_field(u, "control", glyph_factor=1e-3)
viskex.dolfinx.plot_vector_field(z, "adjoint velocity", glyph_factor=1e2)
viskex.dolfinx.plot_scalar_field(b, "adjoint pressure")